use ring::signature::{self, Ed25519KeyPair, KeyPair, UnparsedPublicKey};
use ring::rand::SystemRandom;
use ring::digest;
use data_encoding::BASE64;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

#[derive(Serialize, Deserialize, Clone)]
pub struct KeyPairData {
    pub public_key: String,
    pub private_key: String,
}

#[derive(Debug)]
pub struct CryptoError(pub String);

impl std::fmt::Display for CryptoError {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "Crypto error: {}", self.0)
    }
}

impl std::error::Error for CryptoError {}

pub struct CryptoService {
    key_pairs: HashMap<String, KeyPairData>,
    rng: SystemRandom,
}

impl CryptoService {
    pub fn new() -> Self {
        Self {
            key_pairs: HashMap::new(),
            rng: SystemRandom::new(),
        }
    }

    pub fn generate_keypair(&mut self, peer_id: String) -> Result<KeyPairData, CryptoError> {
        let pkcs8_bytes = Ed25519KeyPair::generate_pkcs8(&self.rng)
            .map_err(|e| CryptoError(format!("Failed to generate keypair: {}", e)))?;

        let key_pair = Ed25519KeyPair::from_pkcs8(pkcs8_bytes.as_ref())
            .map_err(|e| CryptoError(format!("Failed to parse keypair: {}", e)))?;

        let key_data = KeyPairData {
            public_key: BASE64.encode(key_pair.public_key().as_ref()),
            private_key: BASE64.encode(pkcs8_bytes.as_ref()),
        };

        self.key_pairs.insert(peer_id, key_data.clone());
        Ok(key_data)
    }

    pub fn sign_data(&self, peer_id: &str, data: &[u8]) -> Result<String, CryptoError> {
        let key_data = self.key_pairs.get(peer_id)
            .ok_or_else(|| CryptoError(format!("Keypair not found for peer: {}", peer_id)))?;

        let pkcs8_bytes = BASE64.decode(key_data.private_key.as_bytes())
            .map_err(|e| CryptoError(format!("Failed to decode private key: {}", e)))?;

        let key_pair = Ed25519KeyPair::from_pkcs8(&pkcs8_bytes)
            .map_err(|e| CryptoError(format!("Failed to parse private key: {}", e)))?;

        let signature = key_pair.sign(data);
        Ok(BASE64.encode(signature.as_ref()))
    }

    pub fn verify_signature(
        &self, 
        public_key: &str, 
        data: &[u8], 
        signature: &str
    ) -> Result<bool, CryptoError> {
        let public_key_bytes = BASE64.decode(public_key.as_bytes())
            .map_err(|e| CryptoError(format!("Failed to decode public key: {}", e)))?;

        let signature_bytes = BASE64.decode(signature.as_bytes())
            .map_err(|e| CryptoError(format!("Failed to decode signature: {}", e)))?;

        let public_key = UnparsedPublicKey::new(&signature::ED25519, &public_key_bytes);
        
        match public_key.verify(data, &signature_bytes) {
            Ok(()) => Ok(true),
            Err(_) => Ok(false),
        }
    }

    pub fn compute_hash(&self, data: &[u8]) -> String {
        let digest = digest::digest(&digest::SHA256, data);
        BASE64.encode(digest.as_ref())
    }

    pub fn get_public_key(&self, peer_id: &str) -> Option<String> {
        self.key_pairs.get(peer_id).map(|k| k.public_key.clone())
    }
}